#!/usr/bin/python
# -*- coding: UTF-8 -*-
import os
import re

welcome = """
  _______      ____         _                      
 |__   __|    |  _ \       | |                     
    | | ___   | |_) | ___  | |     __ _ _____   _  
    | |/ _ \  |  _ < / _ \ | |    / _` |_  / | | | 
    | | (_) | | |_) |  __/ | |___| (_| |/ /| |_| | 
    |_|\___/  |____/ \___| |______\__,_/___|\__, | 
                                             __/ | 
                                            |___/  
ProtoBuf 生成器，具体格式按照我的demo.txt，石墨笔记直接copy到记事本再运行python脚本， 同行元素以'\\t'分隔
支持数据库的类型转换，如VARCHAR转string，如果是pb类型不做处理，java还不支持，文本格式（注释+字段名+字段类型）
运行：python pbgenerate.py
祝你生活愉快，敲代码不累
"""
now_dir = r'./'
target_dir = 'pb-target'
sql_pb_type_dic = {'VARCHAR': 'string', 'TEXT': 'string',
                   'INT': 'uint32', 'BIGINT': 'string',
                   'TIMESTAMP': 'string'
                   }
java_pb_type_dic = {'int': 'uint32', 'Integer': 'uint32',
                    'long': 'string', 'Long': 'string',
                    'Decimal': 'string', 'String': 'string',
                    'double': 'string', 'float': 'string',
                    'ErrorCode': 'common.errcode.ErrorCode',
                    'PageInfo': 'common.base.Pagination', 
                    'I18nmsg': 'common.base.I18nmsg',
                    'Pagination': 'common.base.Pagination',
                    'bizcode': 'common.enum.Push.MessageBizCodeEnum',
                    'DeviceType': 'common.enum.MsggwEnum.DeviceType'
                    }


def read_file(path):
    """读取文件"""
    file_content = []
    with open(path, encoding='utf-8') as sf:
        try:
            if os.path.isfile(path):
                s = sf.readline()
                while s:
                    s = clear_str(s)
                    s_list = s.split(' ')
                    file_content.append(s_list)
                    s = sf.readline()
                return file_content
        except Exception as e:
            print('[file: {} occur error: {}]'.format(path, e))
            pass


def clear_str(s):
    """格式化字符串"""
    s = s.replace('\t', ' ').replace('\n', ' ').strip(' ')
    return s


def ready_to_write():
    if os.path.exists(target_dir):
        pass
    else:
        os.mkdir(target_dir)


def write_file(file_name, file_content, from_type):
    """写文件， 传入的类型是sql"""
    split_ext = os.path.splitext(file_name)
    message_name = convert_to_message_name(split_ext[0])
    start = 'message ' + message_name + ' {\n'
    with open(file_name, 'w', encoding='utf-8') as wf:
        print(file_content)
        i = 1
        if from_type == 4:
            start = 'enum ' + message_name + ' {\n'
            start += '\t NULL = 0; //空\n'
        for content in file_content:
            if len(content) < 3:
                print('[ERROR]Please use correct format!')
                break
            annotation = content[0].strip(' ')
            is_enum = False
            if from_type == 4:
                is_enum = True
            field_name = convert_field_name(content[1].strip(' '), is_enum)
            field_type = convert_field_type(content, field_name, file_name, from_type)
            start += '\t' + field_type
            start += ' ' + field_name + ' = {};'.format(i)
            i += 1
            start += ' //' + annotation + '\n'
        start += '}\n'
        wf.write(start)


def convert_field_type(content, field_name, file_name, from_type):
    """转换字段类型，如果from_type为4表示enum，不处理类型"""
    field_type = ''
    if from_type == 1:
        field_type = match_type(file_name, field_name, content[2].strip(' '), 1)
    elif from_type == 2:
        field_type = content[2].strip(' ').lower()
    elif from_type == 3:
        field_type = match_type(file_name, field_name, content[2].strip(' '), 3)
    return field_type


def convert_field_name(field_name, is_enum=False):
    """大写转下划线"""
    result = ''
    if field_name:
        if is_enum:
            # 是枚举全大写
            result = field_name.upper()
        else:
            for c in field_name:
                if c.isupper():
                    result += '_' + c.lower()
                else:
                    result += c
    return result


def convert_to_message_name(name):
    """转化为首字母大写"""
    words = name.split("_")
    result = ''
    for word in words:
        result += word.capitalize()
    return result


def match_type(file_name, field_name, field_type, source_type):
    """匹配类型"""
    if source_type == 1:
        for key in sql_pb_type_dic.keys():
            if re.search(key, field_type, re.IGNORECASE):
                # 匹配基础类型
                return sql_pb_type_dic[key]
    if source_type == 3:
        if field_type.startswith('List<'):
            print(field_type)
            # lstrip 实际是按照单个字符来删除的,所以如果像下面这样，List<Long>就会变成ong
            # temp_field_type = field_type.lstrip('List<').rstrip('>')
            temp_field_type = field_type.replace('List<', '').replace('>', '')
            print(temp_field_type)
            sub_type = match_type(file_name, field_name, temp_field_type, source_type)
            return 'repeated ' + sub_type
        for key in java_pb_type_dic.keys():
            if re.search(key, field_type, re.IGNORECASE):
                # 匹配基础类型
                return java_pb_type_dic[key]
        return field_type
    print('[ERROR] ILLEGAL FIELD TYPE IN [FILE] {} [FIELD NAME] {}!'.format(file_name, field_name))
    delete_file()
    exit(0)


def delete_file():
    """删除所有proto结尾的文件"""
    file_list = os.listdir(now_dir)
    for file_path in file_list:
        os.remove(file_path)


def aggregate_data():
    cmd = 'cat * > all.txt'
    os.system(cmd)


def main():
    print(welcome)
    print('***********start to transfer *.txt to proto:***********')
    print('[数据库类型映射关系{}]'.format(sql_pb_type_dic))
    print('[Java类型映射关系{}]'.format(java_pb_type_dic))
    ready_to_write()
    file_list = os.listdir(now_dir)
    abs_path = os.path.abspath(now_dir)
    file_list_abs_path = []
    for file_path in file_list:
        if file_path.endswith('.txt'):
            file_list_abs_path.append(abs_path + '/' + file_path)
    os.chdir(target_dir)
    for file_path in file_list_abs_path:
        print(file_path)
        content_type = int(input("请输入文件%s的字段类型，1是sql，2是pb, 3是java:" % file_path))
        content = read_file(file_path)
        if content_type == 1:
            print('[sql to pb]:')
        elif content_type == 2:
            print('[pb to pb，不转换field_type, 只做大小写转化]:')
        elif content_type == 3:
            print('[java to pb]:')
        else:
            print('未知类型')
            break
        if content is not None:
            file_name = content[0][0].strip(' ').strip('\n')
            content = content[1:]
            if len(content) > 0 and content[0][0] != '':
                if re.search('enum', file_name, re.IGNORECASE):
                    content_type = 4
                write_file(file_name, content, content_type)
    aggregate_data()
    print('************finished*************')


if __name__ == '__main__':
    main()
